#include <stdlib.h>
#include <string.h>
#include "sm2.h"
#include "sm3.h"
#include "asn1.h"

void sm2_compute_z(u8 *z, const sm2_public_key *pubkey, const char *id, u64 idlen) {
    // a ||b || Gx || Gy
    static u8 date[128] =
            {0xff, 0xff, 0xff, 0xfe, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff,
             0xff, 0xff, 0xff, 0xff, 0x00, 0x00, 0x00, 0x00, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xfc,
             0x28, 0xe9, 0xfa, 0x9e, 0x9d, 0x9f, 0x5e, 0x34, 0x4d, 0x5a, 0x9e, 0x4b, 0xcf, 0x65, 0x09, 0xa7,
             0xf3, 0x97, 0x89, 0xf5, 0x15, 0xab, 0x8f, 0x92, 0xdd, 0xbc, 0xbd, 0x41, 0x4d, 0x94, 0x0e, 0x93,
             0x32, 0xc4, 0xae, 0x2c, 0x1f, 0x19, 0x81, 0x19, 0x5f, 0x99, 0x04, 0x46, 0x6a, 0x39, 0xc9, 0x94,
             0x8f, 0xe3, 0x0b, 0xbf, 0xf2, 0x66, 0x0b, 0xe1, 0x71, 0x5a, 0x45, 0x89, 0x33, 0x4c, 0x74, 0xc7,
             0xbc, 0x37, 0x36, 0xa2, 0xf4, 0xf6, 0x77, 0x9c, 0x59, 0xbd, 0xce, 0xe3, 0x6b, 0x69, 0x21, 0x53,
             0xd0, 0xa9, 0x87, 0x7c, 0xc6, 0x2a, 0x47, 0x40, 0x02, 0xdf, 0x32, 0xe5, 0x21, 0x39, 0xf0, 0xa0};

    u32 entl = (idlen << 3);  // bitlen
    u8 *msq = malloc(idlen + 194);

    msq[0] = (entl >> 8) & 0xff;
    msq[1] = entl & 0xff;
    memcpy(msq + 2, id, idlen);
    memcpy(msq + idlen + 2, date, 128);
    memcpy(msq + idlen + 130, pubkey->x, 32);
    memcpy(msq + idlen + 162, pubkey->y, 32);
    sm3(z, msq, entl + 1552);
    free(msq);
}

void sm2_compute_e(u8 *dst, const u8 *z, const u8 *msq, const u64 msqlen) {
    u8 *buf = malloc(msqlen + 32);
    memcpy(buf, z, 32);
    memcpy(buf + 32, msq, msqlen);
    sm3(dst, buf, (msqlen + 32) << 3);
    free(buf);
}

void sm2_keydepair(u8 *output, point p, const u32 mlen) {
    u8 reg[68];
    u8 tmp[32];
    u32 ct = 1;
    u32 n = mlen >> 5;
    fp_to_byt(p.x, reg);
    fp_to_byt(p.y, reg + 32);
    for (; ct < n + 1; ct++) {
        Put(reg + 64, ct);
        sm3(output + ct * 32, reg, 544);
    }
    n = mlen - n;
    Put(reg + 64, ct);
    sm3(tmp, reg, 544);
    memcpy(output + ct * 32 - 32, tmp, n);
}

void sm2_sign_init(u8 *dgst, const sm2_public_key *pubkey, const char *id, u64 idlen, const u8 *msq, u64 msqlen) {
    u8 z[32];
    // compute Za = H ( entl || id || a ||b || Gx || Gy || Px || Py)
    sm2_compute_z(z, pubkey, id, idlen);

    // compute e = H ( Za || M)
    sm2_compute_e(dgst, z, msq, msqlen);
}

void sm2_verify_init(u8 *dgst, const sm2_public_key *pubkey, const char *id, u64 idlen, const u8 *msq, u64 msqlen) {
    u8 z[32];
    // compute Za = H ( entl || id || a ||b || Gx || Gy || Px || Py)
    sm2_compute_z(z, pubkey, id, idlen);

    // compute e = H ( Za || M)
    sm2_compute_e(dgst, z, msq, msqlen);
}

void sm2_sign(sm2_signature *sig, const sm2_private_key *prikey, const u8 *dgst) {

    fp d, k, r, s, t1, t2;
    point kg;

    // retrieve private key
    fp_fm_byt(d, prikey->prikey);

    // compute r, s
    while (1) {
        while (1) {
            // rand k in [1, n - 1]
            fp_rand(k, sm2.n);
            // use for test
            fp_fm_str(k, "59276E27D506861A16680F3AD9C02DCCEF3CC1FA3CDBE4CE6D54B80DEAC1BC21");

            // compute [k]g
            point_mul_fixed(&kg, k);

            //r = e + x_1 (mod n)
            fp_fm_byt(r, dgst);
            fn_add(r, r, kg.x);

            // if r == 0 or r + k == n re-generate k
            fn_add(t1, r, k);
            if (fp_equ(r, zero) == 0 || fp_equ(t1, k) == 0)
                break;
        }
        // s = ((1 + d)^(-1) * (k - r * d)) mod n
        fn_add(t1, d, one);
        fn_inv(t2, t1);
        fn_mul(t1, r, d);
        fn_sub(t1, k, t1);
        fn_mul(s, t1, t2);

        // check s != 0
        if (fp_equ(s, zero) == 0)
            break;
    }

    // return r, s
    fp_to_byt(r, sig->r);
    fp_to_byt(s, sig->s);
}

int sm2_verify(const sm2_signature *sig, const sm2_public_key *publickey, const u8 *dgst){

    fp sigr, sigs, t, r;
    point pubkey, p;

    // retrieve public key
    point_fm_byt(&pubkey, publickey->x, publickey->y);

    // check r′, s′ ∈[1,n-1]
    fp_fm_byt(sigr, sig->r);
    fp_fm_byt(sigs, sig->s);
    if ((fp_tst(sigr) && fp_tst(sigs)) != 1)
        return 0;

    // compute t = r' + s' mod n
    fn_add(t, sigr, sigs);
    if (fp_equ(t, zero))
        return 0;

    // compute s′G + t'p
    point_mutmul(&p, sigs, t, &pubkey);

    // compute r = (e′ + x_1′) mod n
    fp_fm_byt(r, dgst);
    fn_add(r, r, p.x);

    // check r = r'?
    if (fp_equ(r, sigr))
        return 1;
    return 0;
}

void sm2_encrypt(sm2_ciptext *ciphertext, const sm2_public_key *publickey, const char *msq) {

    u32 msqlen = strlen(msq);
    u8 *buf = malloc(64 + msqlen);
    fp k;
    point pubkey, c1, kp;

    // retrieve public key
    point_fm_byt(&pubkey, publickey->x, publickey->y);

    // rand k in [1, n - 1]
    fp_rand(k, sm2.n);
    fp_fm_str(k, "59276E27D506861A16680F3AD9C02DCCEF3CC1FA3CDBE4CE6D54B80DEAC1BC21");

    // compute c1 = [k]G
    point_mul_fixed(&c1, k);
    fp_to_byt(c1.x, ciphertext->x);
    fp_to_byt(c1.y, ciphertext->y);

    // compute c2
    point_mul(&kp, k, &pubkey);
    sm2_keydepair(ciphertext->ctext, kp, msqlen);
    for (u32 i = 0; i < msqlen; i++)
        ciphertext->ctext[i] ^= msq[i];

    // compute hash(x2 || M || y2)
    fp_to_byt(kp.x, buf);
    fp_to_byt(kp.y, buf + 32 + msqlen);
    memcpy(buf + 32, msq, msqlen);
    sm3(ciphertext->hash, buf, (64 + msqlen) << 3);
    ciphertext->clen = msqlen;
    free(buf);
}

int sm2_decrypt(u8 *plaintext, const sm2_private_key *prikey, const sm2_ciptext *ciphertext) {

    u32 msqlen = ciphertext->clen;
    u8 *buf = malloc(64 + msqlen), c3[32];
    fp d;
    point c1, c2;

    // retrieve private key
    fp_fm_byt(d, prikey->prikey);

    // compute c1, c2 = [d]c1

    point_fm_byt(&c1, ciphertext->x, ciphertext->y);
    point_is_on_curve(&c1);
    if (point_is_on_curve(&c1) == 0)
        return -1;
    point_mul(&c2, d, &c1);

    // compute plaintext
    sm2_keydepair(plaintext, c2, msqlen);
    for (u32 i = 0; i < msqlen; i++)
        plaintext[i] ^= ciphertext->ctext[i];

    // compute c3 = hash(x2 || M || y2)
    fp_to_byt(c2.x, buf);
    fp_to_byt(c2.y, buf + 32 + msqlen);
    memcpy(buf + 32, plaintext, msqlen);
    sm3(c3, buf, (64 + msqlen) << 3);
    free(buf);
    // check c3?
    for (int i = 0; i <= 31; i++)
        if (c3[i] != ciphertext->hash[i])
            return 0;
    return 1;
}

void sm2_sig_to_der(sm2_signature sig, u8 **dst, u64 *dstl) {
    u8 buf[70], *_buf = buf;
    u64 blen = 0;
    asn1_integer_to_der(sig.r, 32, &_buf, &blen);
    asn1_integer_to_der(sig.s, 32, &_buf, &blen);
    asn1_sequence_to_der(buf, blen, dst, dstl);
}

void sm2_sig_from_der(sm2_signature *sig, const u8 **src, u64 *srcl) {
    u8 buf[70];
    const u8 *_buf = buf;
    u64 blen = 0, rlen = 0, slen;
    asn1_sequence_from_der(buf, &blen, src, srcl);
    asn1_integer_from_der(sig->r, &rlen, &_buf, &blen);
    asn1_integer_from_der(sig->s, &slen, &_buf, &blen);
}

void sm2_ciptext_to_der(sm2_ciptext ct, u8 **dst, u64 *dstl) {
    u8 *buf = malloc(110 + ct.clen), *_buf = buf;
    u64 blen = 0;
    asn1_integer_to_der(ct.x, 32, &_buf, &blen);
    asn1_integer_to_der(ct.y, 32, &_buf, &blen);
    asn1_octets_to_der(ct.hash, 32, &_buf, &blen);
    asn1_octets_to_der(ct.ctext, ct.clen, &_buf, &blen);
    asn1_sequence_to_der(buf, blen, dst, dstl);
    free(buf);
}

void sm2_ciptext_from_der(sm2_ciptext *ct, const u8 **src, u64 *srcl) {
    u8 buf[300];
    u64 seqlen = 0, genlen = 0, len;
    const u8 *_buf = buf;
    asn1_sequence_from_der(buf, &seqlen, src, srcl);
    asn1_integer_from_der(ct->x, &genlen, &_buf, &seqlen);
    asn1_integer_from_der(ct->y, &genlen, &_buf, &seqlen);
    asn1_octets_from_der(ct->hash, &genlen, &_buf, &seqlen);
    asn1_octets_from_der(ct->ctext, &len, &_buf, &seqlen);
    ct->clen = (u32) len;
}
